package net.aiion.weave.impl.interception.transformation;

import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.Collection;

import net.aiion.weave.asm.Label;
import net.aiion.weave.asm.Opcodes;
import net.aiion.weave.asm.Type;
import net.aiion.weave.spi.bytecode.BytecodeContext;
import net.aiion.weave.spi.bytecode.BytecodeGenerator;
import net.aiion.weave.spi.bytecode.events.AnnotationStartBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.BytecodeEvent;
import net.aiion.weave.spi.bytecode.events.ClassEndBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.ClassStartBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.MethodEndBytecodeEvent;
import net.aiion.weave.spi.bytecode.events.MethodStartBytecodeEvent;
import net.aiion.weave.spi.bytecode.util.BytecodeTransformation;
import net.aiion.weave.spi.interception.Interceptable;
import net.aiion.weave.spi.interception.InterceptionException;
import net.aiion.weave.spi.interception.InterceptionService;

//TODO Add inheritance support (also respect annotations of implemented interfaces).
public class InterceptableTransformation extends BytecodeTransformation
{
    private static final String INTERCEPTION_SUFFIX = "$$Interception";

    private static final String INTERCEPTION_SERVICE_FIELD_NAME = "interceptionService" + INTERCEPTION_SUFFIX;

    private boolean testInMethod;

    private String className;

    private boolean isConstructor;
    private String methodName;
    private String methodDescriptor;
    private boolean methodAnnotated;

    private Collection<BytecodeEvent> suppressedMethodBytecodeEvents;

    public InterceptableTransformation() {
        suppressedMethodBytecodeEvents = new ArrayList<>();
    }

    @Override
    public boolean testRetain(final BytecodeContext aBytecodeContext, final BytecodeEvent aBytecodeEvent) {
        return aBytecodeEvent.getType() == BytecodeEvent.CLASS_START
            && !aBytecodeEvent.asClassStart().containsInterface(Interceptable.class);
    }

    @Override
    public boolean testProcess(final BytecodeContext aBytecodeContext, final BytecodeEvent aBytecodeEvent) {
        switch (aBytecodeEvent.getType()) {
        case BytecodeEvent.METHOD_START:
            testInMethod = true;
            break;
        case BytecodeEvent.METHOD_END:
            testInMethod = false;
            break;
        }
        return aBytecodeEvent.getType() == BytecodeEvent.ANNOTATION_START && testInMethod;
    }

    @Override
    public boolean testSkip(final BytecodeContext aBytecodeContext, final BytecodeEvent aBytecodeEvent) {
        return aBytecodeEvent.getType() == BytecodeEvent.CLASS_START
            && (aBytecodeEvent.asClassStart().isInterface() || aBytecodeEvent.asClassStart().containsInterface(
                Interceptable.class));
    }

    @Override
    public void handleBytecodeEvent(final BytecodeContext aBytecodeContext, final BytecodeEvent aBytecodeEvent) {
        switch (aBytecodeEvent.getType()) {
        case BytecodeEvent.CLASS_START:
            handleClassStartBytecodeEvent(aBytecodeContext, aBytecodeEvent.asClassStart());
            aBytecodeContext.forwardBytecodeEvent(aBytecodeEvent);
            break;
        case BytecodeEvent.ANNOTATION_START:
            handleAnnotationStartBytecodeEvent(aBytecodeContext, aBytecodeEvent.asAnnotationStart());
            break;
        case BytecodeEvent.METHOD_START:
            handleMethodStartBytecodeEvent(aBytecodeContext, aBytecodeEvent.asMethodStart());
            break;
        case BytecodeEvent.METHOD_END:
            handleMethodEndBytecodeEvent(aBytecodeContext, aBytecodeEvent.asMethodEnd());
            break;
        case BytecodeEvent.CLASS_END:
            handleClassEndBytecodeEvent(aBytecodeContext, aBytecodeEvent.asClassEnd());
            aBytecodeContext.forwardBytecodeEvent(aBytecodeEvent);
            break;
        default:
            if (methodName == null || isConstructor) {
                aBytecodeContext.forwardBytecodeEvent(aBytecodeEvent);
            } else {
                suppressedMethodBytecodeEvents.add(aBytecodeEvent);
            }
            break;
        }
    }

    private void handleClassStartBytecodeEvent(final BytecodeContext aBytecodeContext,
        final ClassStartBytecodeEvent aBytecodeEvent)
    {
        className = aBytecodeEvent.getName();
        aBytecodeEvent.addInterface(Interceptable.class);
    }

    private void handleAnnotationStartBytecodeEvent(final BytecodeContext aBytecodeContext,
        final AnnotationStartBytecodeEvent aBytecodeEvent)
    {
        if (methodName == null || isConstructor) {
            aBytecodeContext.forwardBytecodeEvent(aBytecodeEvent);
            return;
        }

        suppressedMethodBytecodeEvents.add(aBytecodeEvent);
        methodAnnotated = true;
    }

    private void handleMethodStartBytecodeEvent(final BytecodeContext aBytecodeContext,
        final MethodStartBytecodeEvent aBytecodeEvent)
    {
        methodName = aBytecodeEvent.getName();
        methodDescriptor = aBytecodeEvent.getDescriptor();
        isConstructor = aBytecodeEvent.isConstructor();

        if (isConstructor) {
            aBytecodeContext.forwardBytecodeEvent(aBytecodeEvent);
            return;
        }

        suppressedMethodBytecodeEvents.add(aBytecodeEvent);
    }

    private void handleMethodEndBytecodeEvent(final BytecodeContext aBytecodeContext,
        final MethodEndBytecodeEvent aBytecodeEvent)
    {
        if (isConstructor) {
            aBytecodeContext.forwardBytecodeEvent(aBytecodeEvent);
            return;
        }

        if (methodAnnotated) {
	        suppressedMethodBytecodeEvents.add(aBytecodeEvent);
	        implementOriginalMethod(aBytecodeContext);
	        implementInterceptedMethod(aBytecodeContext);
	        suppressedMethodBytecodeEvents.clear();
        } else {
        	aBytecodeContext.forwardBytecodeEvents(suppressedMethodBytecodeEvents);
        	suppressedMethodBytecodeEvents.clear();	
        }

        methodName = null;
        methodAnnotated = false;
    }

    private void handleClassEndBytecodeEvent(final BytecodeContext aBytecodeContext,
        final ClassEndBytecodeEvent aBytecodeEvent)
    {
        implementSetInterceptionService(aBytecodeContext);
    }

    private void implementOriginalMethod(final BytecodeContext aBytecodeContext) {
        for (BytecodeEvent bytecodeEvent : suppressedMethodBytecodeEvents) {
            switch (bytecodeEvent.getType()) {
            case BytecodeEvent.METHOD_START:
                MethodStartBytecodeEvent methodStartBytecodeEvent = bytecodeEvent.asMethodStart();
                methodStartBytecodeEvent.setName(methodStartBytecodeEvent.getName() + INTERCEPTION_SUFFIX);
                aBytecodeContext.forwardBytecodeEvent(methodStartBytecodeEvent);
                break;
            default:
                aBytecodeContext.forwardBytecodeEvent(bytecodeEvent);
                break;
            }
        }
    }

    private void implementInterceptedMethod(final BytecodeContext aBytecodeContext) {
        for (BytecodeEvent bytecodeEvent : suppressedMethodBytecodeEvents) {
            aBytecodeContext.forwardBytecodeEvent(bytecodeEvent);
            if (bytecodeEvent.getType() == BytecodeEvent.CODE_START) {
                break;
            }
        }

        final BytecodeGenerator bytecodeGenerator = aBytecodeContext.getBytecodeGenerator();

        final Label startLabel = new Label();
        final Label endLabel = new Label();
        final Label handlerLabel = new Label();

        bytecodeGenerator.markTryCatchBlock(startLabel, endLabel, handlerLabel, "java/lang/NoSuchMethodException");
        bytecodeGenerator.markTryCatchBlock(startLabel, endLabel, handlerLabel, "java/lang/SecurityException");

        bytecodeGenerator.markLabel(startLabel);
        bytecodeGenerator.loadReference(0);
        bytecodeGenerator.loadField(className, INTERCEPTION_SERVICE_FIELD_NAME, InterceptionService.class);

        // arg0
        bytecodeGenerator.loadReference(0);

        // arg1
        bytecodeGenerator.pushConstant(Type.getType("L" + className + ";"));
        bytecodeGenerator.pushConstant(methodName + INTERCEPTION_SUFFIX);

        final Type[] argumentTypes = Type.getArgumentTypes(methodDescriptor);
        final Type returnType = Type.getReturnType(methodDescriptor);

        bytecodeGenerator.pushInteger(argumentTypes.length);
        bytecodeGenerator.newReferenceArray(Class.class);

        for (int i = 0; i < argumentTypes.length; i++) {
            bytecodeGenerator.duplicateReference();
            bytecodeGenerator.pushInteger(i);
            bytecodeGenerator.pushConstant(argumentTypes[i]);
            bytecodeGenerator.storeReferenceArray();
        }

        bytecodeGenerator.invokeVirtual(Class.class, "getDeclaredMethod", Method.class, String.class, Class[].class);

        // arg2
        bytecodeGenerator.pushInteger(argumentTypes.length);
        bytecodeGenerator.newReferenceArray(Object.class);

        for (int i = 0; i < argumentTypes.length; i++) {
            bytecodeGenerator.duplicateReference();
            bytecodeGenerator.pushInteger(i);
            bytecodeGenerator.loadVariable(i, argumentTypes[i]);
            bytecodeGenerator.invokeValueOf(argumentTypes[i]);
            bytecodeGenerator.storeReferenceArray();
        }

        bytecodeGenerator.invokeInterface(InterceptionService.class, "interceptInvocation", Object.class,
            Interceptable.class, Method.class, Object[].class);

        bytecodeGenerator.castChecked(returnType);

        bytecodeGenerator.markLabel(endLabel);
        bytecodeGenerator.returnVariable(returnType);

        bytecodeGenerator.markLabel(handlerLabel);
        bytecodeGenerator.markFrameSame1(0, null, 1, new Object[] {
            "java/lang/Exception"
        });
        bytecodeGenerator.storeReference(2);
        bytecodeGenerator.newReference(InterceptionException.class);
        bytecodeGenerator.duplicateReference();
        bytecodeGenerator.loadReference(2);
        bytecodeGenerator.invokeConstructor(InterceptionException.class, Throwable.class);
        bytecodeGenerator.throwReference();

        bytecodeGenerator.endCode(8, 3);
        bytecodeGenerator.endMethod();
    }

    private void implementSetInterceptionService(final BytecodeContext aBytecodeContext) {
        final BytecodeGenerator bytecodeGenerator = aBytecodeContext.getBytecodeGenerator();
        bytecodeGenerator.startField(Opcodes.ACC_PRIVATE, INTERCEPTION_SERVICE_FIELD_NAME, InterceptionService.class);
        bytecodeGenerator.endField();

        bytecodeGenerator.startMethod(Opcodes.ACC_PUBLIC, "setInterceptionService", void.class,
            InterceptionService.class);
        bytecodeGenerator.startCode();
        bytecodeGenerator.loadReference(0);
        bytecodeGenerator.loadReference(1);
        bytecodeGenerator.storeField(className, INTERCEPTION_SERVICE_FIELD_NAME, InterceptionService.class);
        bytecodeGenerator.returnVoid();
        bytecodeGenerator.endCode(2, 2);
        bytecodeGenerator.endMethod();
    }
}
